import numpy as np
import matplotlib.pyplot as plt

# 生成10x10的随机深度图数组[1,3](@ref)
depth_map = np.random.rand(100, 100) * 10  # 生成0-10米范围的浮点型深度值
# depth_map = np.random.randint(0, 100, (10,10))  # 生成0-100厘米的整型深度值[3](@ref)

def plot_2d_hot_image(depth_map, num_classes=10):
    plt.figure(figsize=(8,6))
    im = plt.imshow(depth_map, 
                    cmap='viridis',  # 颜色映射方案[6](@ref)
                    origin='lower')  # 坐标系方向
    plt.colorbar(im, label='Depth (m)')  # 添加带单位的颜色条[6](@ref)
    plt.title('Random Depth Map')
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.show()

def sudeo_color_image(depth_map):
    import cv2

    norm_depth = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX)
    colored_depth = cv2.applyColorMap(norm_depth.astype(np.uint8), cv2.COLORMAP_JET)

    cv2.imshow('Depth Visualization', colored_depth)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def plot_3d_surface(depth_map):
    from mpl_toolkits.mplot3d import Axes3D

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    x, y = np.meshgrid(range(depth_map.shape[1]), range(depth_map.shape[0]))
    ax.plot_surface(x, y, depth_map, cmap='plasma')

    ax.set_zlabel('Depth (m)')
    plt.show()
    
if __name__ == '__main__':
    plot_2d_hot_image(depth_map)
    sudeo_color_image(depth_map)
    plot_3d_surface(depth_map)